#ifndef WIN32_LEAN_AND_MEAN
#define WIN32_LEAN_AND_MEAN
#endif

#include "Asyncket.h"
#include <MSWSock.h>
#include <Windows.h>

#pragma comment(lib, "ws2_32.lib")

typedef void (*IOCP_CALLBACK)(OVERLAPPED * pOverlapped);

class asyncketInternal
{
private:
	static DWORD WINAPI _workerThread(PVOID lpParameter)
	{
		DWORD size;
		ULONG_PTR key;
		OVERLAPPED * pOverlapped;

		while (GetQueuedCompletionStatus((HANDLE)lpParameter, &size, &key, &pOverlapped, INFINITE)) {
			if (key == 0) {
				break;
			}
			((IOCP_CALLBACK)key)(pOverlapped);
		}
		return 0;
	}

	static void * _getFunctionPtr(SOCKET s, GUID * guid)
	{
		void * p;
		DWORD r;

		if (WSAIoctl(s, SIO_GET_EXTENSION_FUNCTION_POINTER, guid, sizeof(GUID), &p, sizeof(void *), &r, NULL, NULL) == 0) {
			if (r <= sizeof(void *)) {
				return p;
			}
		}
		return NULL;
	}

	static void _defaultCallback(OVERLAPPED * pOverlapped)
	{
	}

public:
	LPFN_ACCEPTEX _AcceptEx;
	LPFN_CONNECTEX _ConnectEx;
	LPFN_DISCONNECTEX _DisconnectEx;

private:
	HANDLE _hHeap;
	HANDLE _hCompletionPort;
	int _workerCount;
	HANDLE * _phWorker;

public:
	asyncketInternal()
		: _workerCount(0),
		_phWorker(NULL)
	{
		SYSTEM_INFO system_info;
		GetSystemInfo(&system_info);
		int workerCount = system_info.dwNumberOfProcessors * 2;
		_hCompletionPort = CreateIoCompletionPort(INVALID_HANDLE_VALUE, 0, 0, 0);
		if (_hCompletionPort == NULL) {
			throw GetLastError();
		}
		_phWorker = new HANDLE[workerCount];
		for (int i = 0; i < workerCount; i++) {
			HANDLE hWorker = CreateThread(NULL, 0, _workerThread, _hCompletionPort, 0, NULL);
			if (hWorker == NULL) {
				throw GetLastError();
			}
			_phWorker[i] = hWorker;
			_workerCount++;
		}
		WSADATA wsa_data;
		WSAStartup(MAKEWORD(2, 2), &wsa_data);
		SOCKET s = socket(AF_UNSPEC, SOCK_STREAM, IPPROTO_TCP);
		GUID GUID_AcceptEx = WSAID_ACCEPTEX;
		_AcceptEx = (LPFN_ACCEPTEX)_getFunctionPtr(s, &GUID_AcceptEx);
		GUID GUID_ConnectEx = WSAID_CONNECTEX;
		_ConnectEx = (LPFN_CONNECTEX)_getFunctionPtr(s, &GUID_ConnectEx);
		GUID GUID_DisconnectEx = WSAID_DISCONNECTEX;
		_DisconnectEx = (LPFN_DISCONNECTEX)_getFunctionPtr(s, &GUID_DisconnectEx);
		closesocket(s);
	}

	~asyncketInternal()
	{
		WSACleanup();
		if (_hCompletionPort != NULL) {
			for (int i = 0; i < _workerCount; i++) {
				PostQueuedCompletionStatus(_hCompletionPort, 0, 0, NULL);
			}
			if (_phWorker != NULL) {
				for (int i = 0, count = 0; i < _workerCount; i += count) {
					count = _workerCount - i;
					if (count >= MAXIMUM_WAIT_OBJECTS) {
						count = MAXIMUM_WAIT_OBJECTS;
					}
					WaitForMultipleObjects(count, &_phWorker[i], TRUE, INFINITE);
				}
				for (int i = 0; i < _workerCount; i++) {
					CloseHandle(_phWorker[i]);
				}
				delete[] _phWorker;
			}
			CloseHandle(_hCompletionPort);
		}
	}

	void BindCompletionPort(HANDLE hDevice, IOCP_CALLBACK callback)
	{
		if (callback == NULL) {
			callback = _defaultCallback;
		}
		if (CreateIoCompletionPort(hDevice, _hCompletionPort, (ULONG_PTR)callback, 0) != _hCompletionPort) {
			throw GetLastError();
		}
	}
};

static asyncketInternal _internal;

static void _asyncketDispatch(OVERLAPPED * pOverlapped)
{
	ASYNCKET_CALLBACK callback = ((ASYNCKET_CONTEXT *)pOverlapped)->Callback;
	if (callback != NULL) {
		callback(pOverlapped);
	}
}

Asyncket::Asyncket()
{
	_s = WSASocket(AF_INET, SOCK_STREAM, IPPROTO_IP, NULL, 0, WSA_FLAG_OVERLAPPED);
	if (_s == INVALID_SOCKET) {
		throw WSAGetLastError();
	}
	_internal.BindCompletionPort((HANDLE)_s, _asyncketDispatch);
}

Asyncket::~Asyncket()
{
	closesocket(_s);
}

SOCKET Asyncket::Object()
{
	return _s;
}

void Asyncket::Bind(struct sockaddr * addr, size_t addrlen)
{
	if (bind(_s, addr, addrlen) == SOCKET_ERROR) {
		throw WSAGetLastError();
	}
}

void Asyncket::Listen(int backlog)
{
	if (listen(_s, backlog) == SOCKET_ERROR) {
		throw WSAGetLastError();
	}
}

void Asyncket::Accept(ASYNCKET_CALLBACK callback, void * data)
{
	ASYNCKET_ACCEPT_CONTEXT * pContext = new ASYNCKET_ACCEPT_CONTEXT;
	ZeroMemory(pContext, sizeof(ASYNCKET_ACCEPT_CONTEXT));
	pContext->Context.pSender = this;
	pContext->Context.Callback = callback;
    pContext->Context.pData = data;
	pContext->Accept = new Asyncket();
	SOCKET a = pContext->Accept->Object();
	DWORD bytes;
	int r = _internal._AcceptEx(_s, a, pContext->SockAddrs, 0, SOCKET_ADDR_SIZE, SOCKET_ADDR_SIZE, &bytes, (LPWSAOVERLAPPED)pContext);
	DWORD e = WSAGetLastError();
	if (r == SOCKET_ERROR && e != WSA_IO_PENDING) {
		throw e;
	}
}

void Asyncket::Connect(struct sockaddr * addr, size_t addrlen, ASYNCKET_CALLBACK callback, void * data)
{
	ASYNCKET_CONTEXT * pContext = new ASYNCKET_CONTEXT;
	ZeroMemory(pContext, sizeof(ASYNCKET_CONTEXT));
	pContext->pSender = this;
	pContext->Callback = callback;
    pContext->pData = data;
	int r = _internal._ConnectEx(_s, addr, addrlen, NULL, 0, NULL, (LPWSAOVERLAPPED)pContext);
	DWORD e = WSAGetLastError();
	if (r == SOCKET_ERROR && e != WSA_IO_PENDING) {
		throw e;
	}
}

void Asyncket::Disconnect(bool reuse, ASYNCKET_CALLBACK callback, void * data)
{
	ASYNCKET_CONTEXT * pContext = new ASYNCKET_CONTEXT;
	ZeroMemory(pContext, sizeof(ASYNCKET_CONTEXT));
	pContext->pSender = this;
	pContext->Callback = callback;
    pContext->pData = data;
	int r = _internal._DisconnectEx(_s, (LPWSAOVERLAPPED)pContext, reuse? TF_REUSE_SOCKET: 0, 0);
	DWORD e = WSAGetLastError();
	if (r == SOCKET_ERROR && e != WSA_IO_PENDING) {
		throw e;
	}
}

void Asyncket::Recv(void * buf, ULONG len, ASYNCKET_CALLBACK callback, void * data)
{
	ASYNCKET_CONTEXT * pContext = new ASYNCKET_CONTEXT;
	ZeroMemory(pContext, sizeof(ASYNCKET_CONTEXT));
	pContext->pSender = this;
	pContext->Callback = callback;
    pContext->pData = data;
	pContext->WSABuf.buf = (CHAR *)buf;
	pContext->WSABuf.len = len;
	DWORD flags = 0;
	DWORD bytes;
	int r = WSARecv(_s, &pContext->WSABuf, 1, &bytes, &flags, (LPWSAOVERLAPPED)pContext, NULL);
	DWORD e = WSAGetLastError();
	if (r == SOCKET_ERROR && e != WSA_IO_PENDING) {
		throw e;
	}
}

void Asyncket::Send(void * buf, ULONG len, ASYNCKET_CALLBACK callback, void * data)
{
	ASYNCKET_CONTEXT * pContext = new ASYNCKET_CONTEXT;
	ZeroMemory(pContext, sizeof(ASYNCKET_CONTEXT));
	pContext->pSender = this;
	pContext->Callback = callback;
    pContext->pData = data;
	pContext->WSABuf.buf = (CHAR *)buf;
	pContext->WSABuf.len = len;
	DWORD bytes;
	int r = WSASend(_s, &pContext->WSABuf, 1, &bytes, 0, (LPWSAOVERLAPPED)pContext, NULL);
	DWORD e = WSAGetLastError();
	if (r == SOCKET_ERROR && e != WSA_IO_PENDING) {
		throw e;
	}
}